import jwt
import datetime
from functools import wraps
from flask import request, jsonify

# JWT配置
SECRET_KEY = "your-secret-key-here"  # 生产环境中应该使用环境变量
ALGORITHM = "HS256"

class JWTManager:
    """JWT令牌管理器"""
    @staticmethod
    def encode_token(payload, expires_delta=None):
        """
        生成JWT令牌
        Args:
            payload (dict): 载荷数据
            expires_delta (datetime.timedelta): 过期时间差
        Returns:
            str: JWT令牌
        """
        if expires_delta:
            expire = datetime.datetime.utcnow() + expires_delta
        else:
            expire = datetime.datetime.utcnow() + datetime.timedelta(hours=1)
        
        payload.update({
            'exp': expire,
            'iat': datetime.datetime.utcnow()
        })
        token = jwt.encode(payload, SECRET_KEY, algorithm=ALGORITHM)
        return token
    
    @staticmethod
    def decode_token(token):
        """
        解码JWT令牌
        
        Args:
            token (str): JWT令牌
            
        Returns:
            dict: 解码后的载荷数据
        """
        try:
            payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
            return payload
        except jwt.ExpiredSignatureError:
            return {"error": "令牌已过期"}
        except jwt.InvalidTokenError:
            return {"error": "无效的令牌"}
    
    @staticmethod
    def generate_access_token(user_id, username):
        """
        生成访问令牌（30分钟）
        
        Args:
            user_id (int): 用户ID
            username (str): 用户名
            
        Returns:
            str: 访问令牌
        """
        payload = {
            'user_id': user_id,
            'username': username,
            'token_type': 'access'
        }
        expires_delta = datetime.timedelta(minutes=30)
        return JWTManager.encode_token(payload, expires_delta)
    
    @staticmethod
    def generate_refresh_token(user_id, username):
        """
        生成刷新令牌（1小时）
        
        Args:
            user_id (int): 用户ID
            username (str): 用户名
            
        Returns:
            str: 刷新令牌
        """
        payload = {
            'user_id': user_id,
            'username': username,
            'token_type': 'refresh'
        }
        expires_delta = datetime.timedelta(hours=1)
        return JWTManager.encode_token(payload, expires_delta)
    
    @staticmethod
    def token_required(f):
        """
        JWT令牌验证装饰器
        
        Args:
            f: 被装饰的函数
            
        Returns:
            function: 装饰后的函数
        """
        @wraps(f)
        def decorated(*args, **kwargs):
            token = None
            
            # 从请求头中获取令牌
            if 'Authorization' in request.headers:
                auth_header = request.headers['Authorization']
                try:
                    token = auth_header.split(" ")[1]  # Bearer <token>
                except IndexError:
                    return jsonify({'message': '令牌格式错误'}), 401
            
            if not token:
                return jsonify({'message': '缺少令牌'}), 401
            
            try:
                data = JWTManager.decode_token(token)
                if 'error' in data:
                    return jsonify({'message': data['error']}), 401
                
                # 将用户信息添加到请求上下文
                request.current_user = data
            except Exception as e:
                return jsonify({'message': '令牌无效'}), 401
            
            return f(*args, **kwargs)
        
        return decorated
